#!/bin/bash

# Docker镜像构建脚本
# 用于构建OA监控系统的Docker镜像

set -e  # 遇到错误时退出

# 配置变量
IMAGE_NAME="kstephenlgz/oa-exam"
TAG="latest"
DOCKERFILE_PATH="."

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查Docker是否运行
check_docker() {
    if ! docker info > /dev/null 2>&1; then
        log_error "Docker is not running. Please start Docker first."
        exit 1
    fi
    log_success "Docker is running"
}

# 检查必要文件
check_files() {
    local required_files=("Dockerfile" "requirements.txt" ".env" "code/main.py")

    for file in "${required_files[@]}"; do
        if [[ ! -f "$file" ]]; then
            log_error "Required file not found: $file"
            exit 1
        fi
    done
    log_success "All required files found"
}

# 清理旧镜像
cleanup_old_images() {
    read -p "Do you want to clean up old images? (y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        log_info "Cleaning up old images..."

        # 删除旧的镜像
        if docker images -q "$IMAGE_NAME:$TAG" > /dev/null 2>&1; then
            docker rmi "$IMAGE_NAME:$TAG" || true
            log_success "Old image removed"
        else
            log_info "No old images to clean up"
        fi

        # 清理悬空镜像
        if [[ $(docker images -f "dangling=true" -q | wc -l) -gt 0 ]]; then
            docker image prune -f
            log_success "Dangling images cleaned up"
        fi
    fi
}

# 构建Docker镜像
build_image() {
    log_info "Building Docker image: $IMAGE_NAME:$TAG"

    # 构建镜像
    if docker build -t "$IMAGE_NAME:$TAG" "$DOCKERFILE_PATH"; then
        log_success "Docker image built successfully: $IMAGE_NAME:$TAG"
        return 0
    else
        log_error "Docker image build failed"
        return 1
    fi
}

# 测试构建的镜像
test_image() {
    log_info "Testing the built image..."

    # 测试基本Python环境
    docker run --rm "$IMAGE_NAME:$TAG" python --version

    # 测试依赖包导入
    local test_script='
import sys
print(f"Python version: {sys.version}")

# 测试依赖包导入
dependencies = [
    ("playwright", "playwright"),
    ("requests", "requests"),
    ("pandas", "pandas"),
    ("pymysql", "pymysql"),
    ("aiohttp", "aiohttp"),
    ("lxml", "lxml"),
    ("dotenv", "python-dotenv"),
    ("loguru", "loguru"),
    ("cryptography", "cryptography")
]

failed_imports = []
for module_name, package_name in dependencies:
    try:
        __import__(module_name)
        print(f"✓ {package_name} import successful")
    except ImportError as e:
        print(f"✗ {package_name} import failed: {e}")
        failed_imports.append(package_name)

# 测试项目模块
try:
    from code.config import USERAGENT
    print("✓ Project modules import successful")
except ImportError as e:
    print(f"✗ Project modules import failed: {e}")
    failed_imports.append("project_modules")

if failed_imports:
    print(f"Failed imports: {failed_imports}")
    sys.exit(1)
else:
    print("✓ All dependency tests passed")
'

    if docker run --rm "$IMAGE_NAME:$TAG" python -c "$test_script"; then
        log_success "Image test passed"
        return 0
    else
        log_error "Image test failed"
        return 1
    fi
}

# 显示镜像信息
show_image_info() {
    log_info "Image information:"
    docker images "$IMAGE_NAME:$TAG"

    log_info "Image size:"
    docker images "$IMAGE_NAME:$TAG" --format "table {{.Repository}}\t{{.Tag}}\t{{.Size}}"
}

# 主函数
main() {
    log_info "Starting Docker image build process..."

    # 检查环境
    check_docker
    check_files

    # 清理旧镜像
    cleanup_old_images

    # 构建镜像
    if build_image; then
        # 测试镜像
        if test_image; then
            show_image_info
            log_success "Docker image build and test completed successfully!"

            echo
            log_info "Next steps:"
            echo "  1. Run the container: ./run_container.sh"
            echo "  2. Or run manually: docker run -d --name oa-exam $IMAGE_NAME:$TAG"
            echo "  3. Check logs: docker logs oa-exam"
        else
            log_error "Image test failed. Please check the build."
            exit 1
        fi
    else
        log_error "Image build failed. Please check the Dockerfile and dependencies."
        exit 1
    fi
}

# 运行主函数
main "$@"